'''compute_ytm.py
Description: Script used to compute bond yields to maturity
'''
import pandas as pd
import numpy as np
import argparse
from numba import jit
import scipy
from scipy import optimize

@jit
def get_face_value_pv(face_value, int_rate, years):
    return face_value / (1 + int_rate)**years
    
@jit
def get_coupon_pv(coupon, int_rate, period, freq):
    pv = coupon / (1 + int_rate/freq)**period
    return pv    
    
class CouponBond:
    
    def get_price(self, coupon, face_value, int_rate, years, freq=1):
        total_coupons_pv = self.get_coupons_pv(coupon, int_rate, years, freq)
        face_value_pv = get_face_value_pv(face_value, int_rate, years)
        result = total_coupons_pv + face_value_pv
        return result
    
    def get_coupons_pv(self, coupon, int_rate, years, freq=1):
        pv = 0
        for period in range(years * freq):
            pv += get_coupon_pv(coupon, int_rate, period+1, freq)
        return pv
    
    def get_ytm(self, bond_price, coupon, years, freq=1, estimate=0.05, face_value=100.):
        get_yield = lambda int_rate: self.get_price(coupon, face_value, int_rate, years, freq) - bond_price
        return optimize.newton(get_yield, estimate)

    def get_ytm_safe(self, bond_price, coupon, years, freq=1, estimate=0.05, face_value=100.):
        try:
            return self.get_ytm(bond_price, coupon, years, freq, estimate, face_value)
        except:
            return np.nan
    
    
if __name__ == "__main__":
    
    # set up parser
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_file', type=str)
    parser.add_argument('--coupon_field', type=str)
    parser.add_argument('--years_to_maturity_field', type=str)
    parser.add_argument('--price_field', type=str)
    parser.add_argument('--out_file', type=str)

    # run parsing
    args = parser.parse_args()
    print("=== Computing yields to maturity ===")
    print("Input file: ", args.input_file)
    print("Coupon field: ", args.coupon_field)
    print("Years to maturity field: ", args.years_to_maturity_field)
    print("Price field: ", args.price_field)
    print("Output file: ", args.out_file)

    ytm_calculator = CouponBond()
    
    data = pd.read_stata(args.input_file)
    data[args.years_to_maturity_field] = data[args.years_to_maturity_field].astype(int)
    
    data['py_ytm'] = data.apply(
        lambda row: ytm_calculator.get_ytm_safe(
            coupon = row[args.coupon_field], 
            bond_price = row[args.price_field], 
            years = row[args.years_to_maturity_field],
            estimate = row[args.coupon_field]/100
        ), axis = 1
    )

    data.to_stata(args.out_file)
